import typing as tp

from . import BaseExecutor
from ..agent import Agent
from .base_aio_process import ShellExecutor as _ShellExecutor
from ..jobs.units import Executor


class ShellExecutor(BaseExecutor):
    def __init__(self, executor: Executor):
        super().__init__(executor)
        self._exe: tp.Optional[_ShellExecutor] = None

    async def run(
            self,
            node: Agent,
            log_callback: tp.Callable[[str, tp.Literal["STDERR", "STDOUT"]], tp.Coroutine] = None
    ) -> int:
        async with _ShellExecutor(
                self._executor.args, timeout=self._executor.timeout, log_callback=log_callback, **self._executor.kwargs
        ) as exe:
            self._exe = exe
            if log_callback:
                await log_callback(F"{self._executor} pid: {self._exe.p.pid}", "STDOUT")
            return await exe.run()

    async def interrupt(self):
        await self._exe.interrupt()
